import torch
import torch.nn as nn


class ELBO():
    """Evidence Lower Bound
    ELBO 是贝叶斯机器学习的损失函数，损失函数的形式为：
        ELBO = Eq(w\theta)(log(w|\theta)-log(w)-log(D|w))
        这是一个期望的形式,不可以直接用梯度下降法去优化,可以用SGVB以及reinforece
        两种方式去估计ELBO.
    Args:
        n_samples ([int]): 采样的数量
    """
    def __init__(self,
                 n_samples=1):
        super().__init__()
        self.n_samples = n_samples

    def sgvb(self, bnn, name, input_data, y_label):
        """Stochastic Gradient Variational Bayes(SGVB) estimator
        Args:
            bnn ([bayesian class instance]): [一个贝叶斯神经网络实例]
            name ([str]): 输出节点的名称
            input_data ([Tensor]): [输入数据]
            y_label ([Tensor]): [标签]
        Returns:
            loss ([Tensor]): [损失]
        """
        loss = 0
        for _ in range(self.n_samples):
            bnn(input_data)
            bnn.node[name].set_observation(y_label)
            loss += bnn.log_prob()
        loss /= self.n_samples
        return loss

    def reinforece(self):
        pass